#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Purpose: Test SDDP-Aquaplan model
# Created: 04/06/2010
# $Id $
# Copyright (c) 2008,2009 by University College London
# Authors:
#  --<>
#
# This program is free software under the GPL (>=v2)
# Read the file COPYING coming with HydroPlatform for details.

import datetime
import numpy
import os
import sys
import unittest

import support

support.init_hydroplatform()

from hydroplatform.model import fields as fd
from hydroplatform.model import nodetypes as nd
from hydroplatform.model import database as db
from hydroplatform.model import network as nt
from hydroplatform.model import units
from hydroplatform.extension.export.common import get_objectfield

from hydroplatform.extension.aquaplan import definition
from hydroplatform.extension.aquaplan import export


def load_sample_network(session):
    units.create_default_units(session)
    # create project
    project = db.Project(u'SDDP-Aquaplan project', 'Test project')
    project.horizon_from = datetime.date(2000,1,1)
    project.horizon_to = datetime.date(2002,12,1)
    project.time_step = units.MONTHLY
    project.projection_params = []
    project.projection_name = u"No projection"
    project.projection_epsg = None
    session.add(project)
    session.commit()
    units.load_units(session, project)
    # register the node definition
    nd.register_definition(session, project, definition)
    
    # create network
    network = db.Network(u'Test network', project)
    session.add(network)
    session.commit()
    
    # create nodes
    nodes = list()
    ndict = dict()

    node_names = [u'Reservoir1', u'Reservoir2', u'Reservoir3', u'Reservoir4', \
                  u"Crop1", u"Crop2", u"Irrigation1", u"Irrigation2", 
                  u"Hydro1", u"Hydro2"]
    for name in node_names:
        ntype = session.query(db.NodeType).filter_by(name=name[:-1]).first() 
        print "Type is %s" % ntype
        node = db.Node(name, project, ntype)
        session.add(node)
        nodes.append(node)
        nlist = ndict.setdefault(ntype, [])
        nlist.append(node)
        
    network.nodes.extend(nodes)
    session.commit()
    
    rnodes = session.query(db.NodeType).filter_by(name=definition.RESERVOIR_NODE).first() 
    fields = [u'smin', u'smax', u'rmin', u'rmax']
    for v in fields:
        fd = get_objectfield(rnodes, v, session)
        for node in ndict[rnodes]:
            # create a random value
            nv = nt.NodeValue(node.id, fd.id, numpy.random.randint(1000))
            session.add(nv)
    ep = get_objectfield(rnodes, u"energy_price", session)
    for node in ndict[rnodes]:
        # create a random value
        nv = nt.NodeValue(node.id, ep.id, numpy.random.random((12,)))
        session.add(nv)    
    ev = get_objectfield(rnodes, u"evaporation", session)
    for node in ndict[rnodes]:
        # create a random value
        nv = nt.NodeValue(node.id, ev.id, numpy.random.random((12,)))
        session.add(nv)            
        
    res = get_objectfield(rnodes, u"reservoir", session)
    for node in ndict[rnodes]:
        # create a random value
        nv = nt.NodeValue(node.id, res.id, numpy.random.random((numpy.random.randint(1,5),3)))
        session.add(nv)          
    session.commit()

    cnodes = session.query(db.NodeType).filter_by(name=definition.CROP_NODE).first() 
    fields = [u'yield', u'crop_price', u'first_irrigation_month', u'ky']
    for v in fields:
        fd = get_objectfield(cnodes, v, session)
        for node in ndict[cnodes]:
            # create a random value
            nv = nt.NodeValue(node.id, fd.id, numpy.random.randint(1000))
            session.add(nv)    
    wn = get_objectfield(cnodes, u"water_need", session)
    for node in ndict[cnodes]:
        # create a random value
        nv = nt.NodeValue(node.id, wn.id, numpy.random.random((12,)))
        session.add(nv)    
    session.commit()            
    
            
    lnktype = session.query(db.LinkType).filter_by(name=definition.IRRIGATION_LINK).first() 
    print lnktype
    
    inodes = session.query(db.NodeType).filter_by(name=definition.IRRIGATION_NODE).first() 
    fields = [u'efficiency', u'surface', u'return_efficiency']
    for v in fields:
        fd = get_objectfield(inodes, v, session)
        for node in ndict[inodes]:
            # create a random value
            nv = nt.NodeValue(node.id, fd.id, numpy.random.randint(1000))
            session.add(nv)    
    crop = get_objectfield(inodes, u"crop", session)
    for node in ndict[inodes]:
        # create a random value
        nv = nt.NodeValue(node.id, crop.id, ndict[cnodes][numpy.random.randint(2)].id)
        from_link = nt.Link(project, ndict[rnodes][0], node, lnktype)
        return_link = nt.Link(project, node, ndict[rnodes][-1], lnktype)
        session.add(nv)    
        session.add(from_link)
        session.add(return_link)
    session.commit()                        
    
    
    hnodes = session.query(db.NodeType).filter_by(name=definition.HYDRO_NODE).first() 
    hydrology = get_objectfield(hnodes, u"hydrology", session)
    for node in ndict[hnodes]:
        nv = nt.NodeValue(node.id, hydrology.id, numpy.random.random(len(hydrology.attribute_type.dates)))
        session.add(nv)    
    session.commit()                        
    
    return project
    
    

class TestModel(unittest.TestCase):

    def setUp(self):
        """
        """
        self.session = support.get_dbsession()
    
    def test_model(self):
        self.project = load_sample_network(self.session)


class TestExport(unittest.TestCase):
    
    def setUp(self):
        self.session = support.get_dbsession()
        self.project = load_sample_network(self.session)
        self.basedir = os.path.dirname(__file__)
        self.outputfile = os.path.join(self.basedir, "sddp_aquaplan_test.xls")
    
    def test_export(self):
        network = self.project.networks[0]
        exporter = export.AquaplanExporter()
        exporter.export(network, self.outputfile, self.session)
            
    def tearDown(self):
        os.remove(self.outputfile)
        import glob
        for csvfile in glob.glob(os.path.join(self.basedir, "*.csv")):
            print csvfile
            os.remove(csvfile)

if __name__=='__main__':
    unittest.main()